use chumsky::prelude::*;

/// Defines the tokens for the MATLAB language.
/// This enum represents all the fundamental building blocks (lexemes) of MATLAB.
#[derive(Debug, PartialEq, Clone)]
pub enum Token {
    /// Reserved keywords like `if`, `for`, `function`, etc.
    Keyword(String),
    /// Identifiers for variables, functions, and so on.
    Identifier(String),
    /// Numeric literals, supporting integers, floats, and scientific notation.
    Number(f64),
    /// String literals enclosed in single or double quotes.
    String(String),
    /// Operators such as `+`, `-`, `*`, `/`, `==`, etc.
    Op(String),
    /// Punctuation marks like `[`, `]`, `(`, `)`, etc.
    Punct(char),
    /// A single line comment starting with `%`.
    Comment(String),
    /// The semicolon, which suppresses output.
    SemiColon,
}

use std::fmt;

impl fmt::Display for Token {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Token::Keyword(s) => write!(f, "Keyword({})", s),
            Token::Identifier(s) => write!(f, "Identifier({})", s),
            Token::Number(n) => write!(f, "Number({})", n),
            Token::String(s) => write!(f, "String(\"{}\")", s),
            Token::Op(s) => write!(f, "Op({})", s),
            Token::Punct(c) => write!(f, "Punct({})", c),
            Token::Comment(s) => write!(f, "Comment(%){}", s),
            Token::SemiColon => write!(f, "SemiColon(;)"),
        }
    }
}


/// Constructs the lexer for a MATLAB-like language.
pub fn lexer<'a>() -> impl Parser<'a, &'a str, Vec<Token>, extra::Err<Simple<'a, char>>> {
    let comment = just("%")
        .ignore_then(none_of("\n").repeated())
        .to_slice()
        .map(|s: &str| Token::Comment(s.to_string()));

    let op = choice((
        just("=="), just("~="), just(">="), just("<="),
        just("&&"), just("||"),
        just("+"), just("-"), just("*"), just("/"), just("\\"), just("^"),
        just("="), just(">"), just("<"), just("&"), just("|"), just("~"),
    )).map(|s| Token::Op(s.to_string()));

    let punct = one_of("[](){},:").map(Token::Punct);

    let semicolon = just(";").to(Token::SemiColon);

    let num = text::digits(10)
        .then(just(".").then(text::digits(10)).or_not())
        .then(
            just("e").or(just("E"))
                .then(just("+").or(just("-")).or_not())
                .then(text::digits(10))
                .or_not()
        )
        .to_slice()
        .map(|s: &str| Token::Number(s.parse().unwrap()));

    let string_double = just("\"")
        .ignore_then(none_of("\"").repeated())
        .then_ignore(just("\""))
        .to_slice()
        .map(|s:&str| Token::String(s.to_string()));

    let string_single = just("'")
        .ignore_then(none_of("'").repeated())
        .then_ignore(just("'"))
        .to_slice()
        .map(|s:&str| Token::String(s.to_string()));

    let ident = text::ascii::ident()
        .to_slice()
        .map(|s: &str| {
            match s {
                "if" | "else" | "elseif" | "end" | "for" | "while" | "function" | "return" | "break" | "continue" => Token::Keyword(s.to_string()),
                _ => Token::Identifier(s.to_string()),
            }
        });

    let token = choice((
        comment,
        op,
        punct,
        semicolon,
        num,
        string_double,
        string_single,
        ident,
    ))
    .padded();

    token.repeated().collect().then_ignore(end())
}


#[cfg(test)]
mod tests{
    use chumsky::Parser;

    use crate::parser::lexer::lexer;

#[test]
fn main() {
    // --- Test Cases ---
    let source_code = r#"
        % This is a simple comment.
        function [output1, output2] = my_func(input_1)
            % A nested comment with various elements.
            for i = 1:10
                if i > 5.5
                    % String literals can use both single and double quotes.
                    disp("i is greater than 5.");
                    disp('i value: ');
                    
                    % This shows matrix-like syntax and different number formats.
                    A = [1 2; 3 4];
                    B = [5; 6];
                    C = [7.89e-2, 10];
                end
            end
            
            % An example of logical operators.
            if (a == b) && (c ~= d)
                disp("A and B are equal.");
            end
            
            % Another example with an empty matrix.
            D = [];
        end
    "#;

    println!("--- Parsing MATLAB Source Code ---\n");
    match lexer().parse(source_code).into_result() {
        Ok(tokens) => {
            println!("Lexing successful. Found {} tokens.\n", tokens.len());
            for token in tokens {
                println!("{:?}", token);
            }
        }
        Err(errors) => {
            eprintln!("Lexing failed.");
            for e in errors {
                eprintln!("Error: {:?}", e);
            }
        }
    }
}
}
